#!/bin/bash

# FinBERT API 测试脚本
# 使用方法: ./test_api.sh [API_BASE_URL]
# 默认 API 地址: http://localhost:8000

API_BASE_URL="${1:-http://localhost:8000}"

# 颜色定义
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}FinBERT API 测试脚本${NC}"
echo -e "${BLUE}API 地址: ${API_BASE_URL}${NC}"
echo -e "${BLUE}========================================${NC}\n"

# 测试 1: 健康检查
echo -e "${YELLOW}[测试 1] 健康检查${NC}"
echo -e "${GREEN}GET ${API_BASE_URL}/health${NC}"
curl -s -X GET "${API_BASE_URL}/health" | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 2: 获取标签列表
echo -e "${YELLOW}[测试 2] 获取模型支持的情感标签${NC}"
echo -e "${GREEN}GET ${API_BASE_URL}/api/v1/labels${NC}"
curl -s -X GET "${API_BASE_URL}/v1/labels" | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 3: 单个文本情感分析（基础）
echo -e "${YELLOW}[测试 3] 单个文本情感分析（基础）${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": "The company'\''s quarterly earnings exceeded all expectations, sending the stock soaring."
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 4: 单个文本情感分析（返回所有分数）
echo -e "${YELLOW}[测试 4] 单个文本情感分析（返回所有分数）${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": "Investors are concerned about the potential impact of new regulations on the tech sector.",
    "return_all_scores": true
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 5: 批量文本情感分析
echo -e "${YELLOW}[测试 5] 批量文本情感分析${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": [
      "The company'\''s quarterly earnings exceeded all expectations, sending the stock soaring.",
      "Investors are concerned about the potential impact of new regulations on the tech sector.",
      "The market remained stable as traders awaited the latest inflation report."
    ],
    "return_all_scores": false
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 6: 批量文本情感分析（返回所有分数）
echo -e "${YELLOW}[测试 6] 批量文本情感分析（返回所有分数）${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": [
      "The company'\''s quarterly earnings exceeded all expectations, sending the stock soaring.",
      "The market remained stable as traders awaited the latest inflation report."
    ],
    "return_all_scores": true
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 7: 使用 classify 端点（别名）
echo -e "${YELLOW}[测试 7] 使用 classify 端点（别名）${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/classify${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/classify" \
  -H "Content-Type: application/json" \
  -d '{
    "text": "The market remained stable as traders awaited the latest inflation report.",
    "return_all_scores": true,
    "max_length": 512
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 8: 中文文本测试（如果模型支持）
echo -e "${YELLOW}[测试 8] 中文文本测试${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": "公司季度收益超出所有预期，股价飙升。",
    "return_all_scores": true
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 9: 自定义 max_length 参数
echo -e "${YELLOW}[测试 9] 自定义 max_length 参数${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": "The company announced record-breaking profits this quarter, with revenue growth of 50% year-over-year, significantly outperforming market expectations.",
    "max_length": 256,
    "return_all_scores": false
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 10: Prometheus 指标
echo -e "${YELLOW}[测试 10] Prometheus 指标${NC}"
echo -e "${GREEN}GET ${API_BASE_URL}/metrics${NC}"
curl -s -X GET "${API_BASE_URL}/metrics" | head -20
echo -e "\n"

# 测试 11: 错误处理测试（空文本）
echo -e "${YELLOW}[测试 11] 错误处理测试（空文本）${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": ""
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

# 测试 12: 长文本测试
echo -e "${YELLOW}[测试 12] 长文本测试${NC}"
echo -e "${GREEN}POST ${API_BASE_URL}/api/v1/sentiment${NC}"
curl -s -X POST "${API_BASE_URL}/api/v1/sentiment" \
  -H "Content-Type: application/json" \
  -d '{
    "text": "The financial markets experienced significant volatility today as investors reacted to the latest economic data. The Federal Reserve'\''s announcement regarding interest rates has created uncertainty in the market. Analysts are divided on whether this will lead to a bullish or bearish trend in the coming weeks. Technology stocks, in particular, have shown mixed signals, with some companies reporting strong earnings while others face regulatory challenges.",
    "return_all_scores": true
  }' | jq '.' || echo "响应不是有效的 JSON"
echo -e "\n"

echo -e "${BLUE}========================================${NC}"
echo -e "${GREEN}所有测试完成！${NC}"
echo -e "${BLUE}========================================${NC}"

